import os
import csv
from ai_trans import translate_sentences

# 设置 OpenAI API 配置
os.environ['OPENAI_API_KEY'] = 'sk-NscqaCD1PfVm7soEF3C3E6297bE14d7fB595Be8f17F39aFf'
os.environ['OPENAI_API_BASE'] = 'https://aiapi.magong.site/v1'

def read_csv(file_path):
    """读取CSV文件并返回数据列表"""
    encodings = ['utf-8', 'latin-1', 'gbk']
    for encoding in encodings:
        try:
            with open(file_path, mode='r', encoding=encoding) as file:
                reader = csv.reader(file)
                data = [row for row in reader]
            return data
        except UnicodeDecodeError:
            continue
        except Exception as e:
            print(f"Error reading CSV file with encoding {encoding}: {e}")
            return []
    print("Failed to read CSV file with any of the specified encodings.")
    return []


def insert_empty_column(data, column_index):
    """在指定列之后插入一个空列"""
    for row in data:
        row.insert(column_index + 1, '')  # 插入在目标列后面
    return data





def translate_column(data, column_index, start_row=0, target_language='zh'):
    """
    翻译指定列的文本，并将结果保存到下一列
    :param data: CSV数据列表
    :param column_index: 需要翻译的列索引
    :param start_row: 开始翻译的行索引，默认为第0行
    :param target_language: 目标语言，默认为中文
    :return: 翻译后的数据列表
    """
    sentences_to_translate = [row[column_index] for i, row in enumerate(data[start_row:], start=start_row) if len(row) > column_index]
    
    if sentences_to_translate:
        try:
            print(f"Translating {len(sentences_to_translate)} sentences to {target_language}.")
            print(f"Input sentences: {sentences_to_translate}")
            translated_output = translate_sentences(sentences_to_translate, target_language, api_key='YOUR_API_KEY')
            print(f"Translated output: {translated_output}")
            translations = translated_output.get("translations", [])
            
            for i, row in enumerate(data[start_row:], start=start_row):
                if len(row) > column_index and i - start_row < len(translations):
                    row[column_index + 1] = translations[i - start_row]  # 写入到下一列
                    print(f"Translated row {i}: {row[column_index]} -> {row[column_index + 1]}")
        except Exception as e:
            print(f"Error translating rows: {e}")
    
    return data


def save_csv(data, file_path):
    """将数据保存为CSV文件"""
    try:
        with open(file_path, mode='w', newline='', encoding='utf-8') as file:
            writer = csv.writer(file)
            writer.writerows(data)
    except Exception as e:
        print(f"Error saving CSV file: {e}")


def process_csv(input_file, output_file, column_index, start_row=0, target_language='zh'):
    """处理CSV文件的主要函数"""
    data = read_csv(input_file)

    # 插入空列
    data = insert_empty_column(data, column_index)

    # 翻译第二列的文本并保存到下一列
    data = translate_column(data, column_index, start_row, target_language)

    # 保存为新文件
    save_csv(data, output_file)


if __name__ == "__main__":
    input_file = "18.JP家具 - 小的.csv"
    output_file = "18.JP家具 - 小的_processed.csv"
    column_index = 1  # 插入空列的列索引（第2列）
    start_row = 2  # 从第2行开始翻译（通常第0行是标题）

    process_csv(input_file, output_file, column_index, start_row)